#include "Decryptor.h"
using namespace std;

int errno;

Decryptor::Decryptor(const char* key)
{
	strncpy(this->key,key,512);

	register_algs();
}

Decryptor::~Decryptor()
{

}

void Decryptor::register_algs(void)
{
   register_cipher (&aes_desc);

   if (register_hash(&sha256_desc) == -1) {
      printf("Error registering SHA256\n");
      exit(-1);
   }

   if (register_prng(&yarrow_desc) == -1) {
      printf("Error registering yarrow PRNG\n");
      exit(-1);
   }

   if (register_prng(&sprng_desc) == -1) {
      printf("Error registering sprng PRNG\n");
      exit(-1);
   }
}

/* use tomcrypt functions to decrypt the specified file and return the contents in a string */
/* using the password provided in the constructor */
string Decryptor::decrypt(const char* filename)
{
   unsigned char plaintext[512];
   unsigned char key[MAXBLOCKSIZE], IV[MAXBLOCKSIZE];
   unsigned char inbuf[512]; /* i/o block size */
   unsigned long outlen, y, ivsize, decrypt;
   symmetric_CTR ctr;
   int cipher_idx, hash_idx, ks;
   ifstream fdin;
   ostringstream fdout(ostringstream::out);
   string result;

   fdin.open(filename,ios::binary | ios::in);

   cipher_idx = find_cipher("aes");
   if (cipher_idx == -1) {
      printf("Invalid cipher.\n");
      exit(-1);
   }

   hash_idx = find_hash("sha256");
   if (hash_idx == -1) {
      printf("SHA256 not found...?\n");
      exit(-1);
   }

   ivsize = cipher_descriptor[cipher_idx].block_length;
   ks = hash_descriptor[hash_idx].hashsize;
   if (cipher_descriptor[cipher_idx].keysize(&ks) != CRYPT_OK) {
      printf("Invalid keysize???\n");
      exit(-1);
   }

   outlen = sizeof(key);
   if ((errno = hash_memory(hash_idx,(unsigned char*)this->key,strlen((char *)this->key),key,&outlen)) != CRYPT_OK) {
      printf("Error hashing key: %s\n", error_to_string(errno));
      exit(-1);
   } //else { printf("#%s#\n",(char *)this->key); printf("#%s#\n",(char *)key); }

   decrypt = 1;


   if (decrypt) {
        printf("Decrypting...\n");
      /* Need to read in IV */
      if (fdin.read((char *)IV,ivsize).fail() || fdin.eof() || (unsigned long)fdin.gcount() != ivsize) {
         printf("Error reading IV from input.\n");
         exit(-1);
      }

      if ((errno = ctr_start(cipher_idx,IV,key,ks,0,CTR_COUNTER_LITTLE_ENDIAN,&ctr)) != CRYPT_OK) {
         printf("ctr_start error: %s\n",error_to_string(errno));
         exit(-1);
      }

      /* IV done */
      do {
         y = fdin.read((char *)inbuf,sizeof(inbuf)).gcount();
		 //cout << filename << ":" << y << endl;
         if ((errno = ctr_decrypt(inbuf,plaintext,y,&ctr)) != CRYPT_OK) {
            printf("ctr_decrypt error: %s\n", error_to_string(errno));
            exit(-1);
         }

		 //cout << filename << ":" << inbuf << ":" << plaintext << endl;
         if (fdout.write((char *)plaintext,y).bad()) {
            printf("Error writing to file.\n");
            exit(-1);
         }
      } while (y == sizeof(inbuf));
      fdin.close();
      result = fdout.str();

      printf("Decryption finished successfully!\n");

      return result;
   }

   printf("Uh-oh this should not trigger...\n");

   return "";
}

/* thin wrapper to avoid having to add stringToWString on every call to decrypt */
wstring Decryptor::decrypt_w(const char* filename)
{
    string temp = this->decrypt(filename);
    return this->stringToWString(temp);
}

/* Convert the contents of string, assumed to be UTF-8, to wchar_t in a wstring */
wstring Decryptor::stringToWString(const string& s)
{
    int result_length;
    //DWORD error = 0;
    wchar_t wtemp[s.length()];
	wstring temp;
#ifdef _WIN32
    result_length = MultiByteToWideChar(65001,NULL,s.c_str(),-1,wtemp,s.length());
#else
    setlocale(LC_ALL, "en_US.UTF-8");
	result_length = mbstowcs(wtemp,s.c_str(),s.length());
#endif
	temp = wtemp;
	return temp;
}

/* Convert the contents of wstring to UTF-8 in string */
string Decryptor::wstringToString(const wstring& s)
{
    int result_length;
    char ctemp[s.length()*2];
	string temp;
#ifdef _WIN32
    result_length = WideCharToMultiByte(65001,NULL,s.c_str(),-1,ctemp,s.length()*4,NULL,NULL);
#else
    setlocale(LC_ALL, "en_US.UTF-8");
	result_length = wcstombs(ctemp,s.c_str(),s.length()*2);
    if (result_length == -1)
    {
        perror("wcstombs failed");
        exit(1);
    }

#endif
	temp = ctemp;
	return temp;
}
